from faker import Faker
from unidecode import unidecode
from os.path import basename, getsize
from zipfile import ZipFile, ZIP_LZMA
from rich.progress import Progress
from rich import print
# import outside modules
from common.core import human_readable

LOCALES = {
    'English': 'en_US',
    'French': 'fr_FR',
    'German': 'de_DE',
    'Czech': 'cs_CZ',
    'Chinese': 'zh_CN',
    'Japanese': 'ja_JP',
    'Russian': 'ru_RU',
    'Spanish': 'es_ES',
    'Italian': 'it_IT'
}

ZIP_PATH = './name_classification.zip'
# TOTAL_NAMES = 9000000
# NUM_COUNTRIES = TOTAL_NAMES // len(LOCALES)
NUM_COUNTRIES = 1000000

with ZipFile(ZIP_PATH, 'w', ZIP_LZMA) as zip_f, Progress() as pbar:
    txt_path = ZIP_PATH.replace('.zip', '.txt')
    with zip_f.open(basename(txt_path), 'w') as f:
        for country, locale in LOCALES.items():
            fake = Faker(locale=locale)
            task = pbar.add_task(f'[cyan]{country}[/]', total=NUM_COUNTRIES)
            for _ in range(NUM_COUNTRIES):
                name = fake.first_name()
                ascii_name = unidecode(name).strip()
                line = f'{ascii_name}\t{country}\n'.encode('utf-8')
                f.write(line)
                pbar.update(task, advance=1)

print(f'[bold green]data file saved: `{ZIP_PATH}`[/]')
file_size = human_readable(getsize(ZIP_PATH))
print(f'[bold green]zip file size: {file_size}[/]')
